import pymysql
from pymysql.cursors import DictCursor

class MySQLDatabase:
    def __init__(self, db_config):
        """初始化数据库连接"""
        self.db_config = db_config
        self.connection = None
        self.connect()

    def connect(self):
        """建立数据库连接"""
        try:
            self.connection = pymysql.connect(**self.db_config)
        except pymysql.MySQLError as e:
            raise

    def execute_query(self, query, params=None):
        """执行查询操作"""
        try:
            with self.connection.cursor(DictCursor) as cursor:
                cursor.execute(query, params or ())
                return cursor.fetchall()
        except pymysql.MySQLError as e:
            return []

    def execute_non_query(self, query, params=None):
        """执行非查询操作（INSERT/UPDATE/DELETE）"""
        try:
            with self.connection.cursor() as cursor:
                cursor.execute(query, params or ())
                self.connection.commit()
                return cursor.rowcount
        except pymysql.MySQLError as e:
            self.connection.rollback()
            return 0

    def insert(self, table, data):
        """通用插入方法"""
        columns = ', '.join(data.keys())
        placeholders = ', '.join(['%s'] * len(data))
        query = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
        return self.execute_non_query(query, tuple(data.values()))

    def update(self, table, data, where_clause, where_params):
        """通用更新方法"""
        set_clause = ', '.join([f"{key} = %s" for key in data.keys()])
        query = f"UPDATE {table} SET {set_clause} WHERE {where_clause}"
        params = tuple(data.values()) + tuple(where_params)
        return self.execute_non_query(query, params)

    def delete(self, table, where_clause, where_params):
        """通用删除方法"""
        query = f"DELETE FROM {table} WHERE {where_clause}"
        return self.execute_non_query(query, tuple(where_params))

    def fetch_one(self, query, params=None):
        """获取单条记录"""
        try:
            with self.connection.cursor(DictCursor) as cursor:
                cursor.execute(query, params or ())
                return cursor.fetchone()
        except pymysql.MySQLError as e:
            return None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def close(self):
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()